/**
 * Copyright (C) 2021 - 2031 O-Cubes Co., Ltd.
 */

/****************************************************************
 *  @file    isp.c
 *  @brief   ISP Program driver
 *  @version v1.0
 *  @date    03. Apr. 2023
 *  @author  gonght
 ****************************************************************/

#include "isp.h"
#include "printf.h"
#include "error.h"
#include "sfud.h"
#include "dw_uart.h"
#include "device_uart.h"
#include "image_header.h"
#include "syscounter.h"

uint8_t g_flash_tmp_buf[ISP_FLASH_PKG_HDR_SIZE];

#define ISP_CHECK_EXP(EXPR, ERROR_CODE) CHECK_EXP(EXPR, ercd, ERROR_CODE, error_exit)

static void isp_cmd_get_info_handler(void);
static void isp_cmd_set_baud_handler(void);
static void isp_cmd_read_mem_handler(void);
static void isp_cmd_go_handler(void);
static void isp_cmd_rst_handler(void);
static void isp_cmd_write_mem_handler(void);

#ifdef CONFIG_SDK_QSPIFLASH
static void isp_cmd_read_flash_handler(void);
static void isp_cmd_write_flash_handler(void);
static void isp_cmd_erase_flash_handler(void);
static void isp_cmd_upgrade_handler(void);
static void isp_cmd_get_boot_part_handler(void);
static void isp_cmd_switch_boot_part_handler(void);
#endif

static const isp_cmd_func_set glb_isp_func_set[] = {
	{ ISP_CMD_GET_INFO,           isp_cmd_get_info_handler          },
	{ ISP_CMD_SET_BAUD,           isp_cmd_set_baud_handler          },
	{ ISP_CMD_READ_MEMORY,        isp_cmd_read_mem_handler          },
	{ ISP_CMD_RST,                isp_cmd_rst_handler               },
	{ ISP_CMD_WRITE_MEMORY,       isp_cmd_write_mem_handler         },
#ifdef CONFIG_SDK_QSPIFLASH
	{ ISP_CMD_READ_FLASH,         isp_cmd_read_flash_handler        },
	{ ISP_CMD_WRITE_FLASH,        isp_cmd_write_flash_handler       },
	{ ISP_CMD_ERASE_FLASH,        isp_cmd_erase_flash_handler       },
	{ ISP_CMD_UPGRADE,            isp_cmd_upgrade_handler           },
	{ ISP_CMD_GET_BOOT_PART,      isp_cmd_get_boot_part_handler     },
	{ ISP_CMD_SWITCH_BOOT_PART,   isp_cmd_switch_boot_part_handler  },
#endif
	{ ISP_CMD_STOP,               NULL  },
};

int32_t isp_uart_putc(const unsigned char c)
{
	int32_t ercd = E_OK;

	DEV_UART_PTR uart_ptr = uart_get_dev(ISP_DFT_UART_ID);

	ISP_CHECK_EXP(uart_ptr != NULL, E_OBJ);

	ercd = uart_ptr->uart_write(&c, 1);
	ISP_CHECK_EXP(ercd == 1, E_SYS);

error_exit:
	return ercd;
}

unsigned char isp_uart_getc(void)
{
	int32_t ercd = E_OK;
	char c;

	DEV_UART_PTR uart_ptr = uart_get_dev(ISP_DFT_UART_ID);

	ISP_CHECK_EXP(uart_ptr != NULL, E_OBJ);

	ercd = uart_ptr->uart_read(&c, 1);
	ISP_CHECK_EXP(ercd == 1, E_SYS);

	return c;

error_exit:
	return ercd;
}

int32_t isp_uart_puts(const void *str, uint32_t len)
{
	int32_t ercd = E_OK;

	DEV_UART_PTR uart_ptr = uart_get_dev(ISP_DFT_UART_ID);

	ISP_CHECK_EXP(uart_ptr != NULL, E_OBJ);
	ISP_CHECK_EXP(str != NULL, E_PAR);

	ercd = uart_ptr->uart_write(str, len);
	ISP_CHECK_EXP(ercd == len, E_SYS);

error_exit:
	return ercd;
}

int32_t isp_uart_gets(void *str, uint32_t len)
{
	int32_t ercd = E_OK;

	DEV_UART_PTR uart_ptr = uart_get_dev(ISP_DFT_UART_ID);

	ISP_CHECK_EXP(uart_ptr != NULL, E_OBJ);
	ISP_CHECK_EXP(str != NULL, E_PAR);

	ercd = uart_ptr->uart_read(str, len);
	ISP_CHECK_EXP(ercd == len, E_SYS);

error_exit:
	return ercd;
}

int32_t isp_set_uart_baud(uint32_t baud)
{
	int32_t ercd = E_OK;

	DEV_UART_PTR uart_ptr = uart_get_dev(ISP_DFT_UART_ID);

	ISP_CHECK_EXP(uart_ptr != NULL, E_OBJ);

	ercd = uart_ptr->uart_control(UART_CMD_SET_BAUD, CONV2VOID(baud));
	ISP_CHECK_EXP(ercd == E_OK, E_SYS);

error_exit:
	return ercd;
}

static void isp_send_ack(void)
{
	isp_uart_putc(ISP_CMD_ACK);
}

static void isp_send_nack(uint8_t error)
{
	isp_uart_putc(ISP_CMD_NACK);
	isp_uart_putc(error);
}

static uint8_t isp_calc_checksum(uint8_t *data, uint32_t len)
{
	uint32_t i;
	uint8_t cs = 0;

	for (i = 0; i < len; i++) {
		cs ^= data[i];
	}

	return cs;
}

static int32_t isp_cmp_checksum(uint8_t cs, uint8_t *data, uint32_t len)
{
	if (cs == isp_calc_checksum(data, len))
		return E_OK;
	else
		return E_FAILURE;
}

static int32_t isp_upgrade_pkg_check(uint32_t bak_partition)
{
	int32_t result = E_OK;
	struct pkg_header *pkg_hdr = (struct pkg_header *)PKG_HDR_FLASH_BASE;

#ifndef SPIFLASH_XIP_EN
	sfud_xip_en();
#endif

	result = pkg_header_check(pkg_hdr);
	if (result) {
		return E_FAILURE;
	}

	result = subimg_data_check((uint8_t *)MEM_MAP_FLASH_BASE_ADDR + ntohl(pkg_hdr->subimg_hdr[bak_partition].part_base),
			ntohl(pkg_hdr->subimg_hdr[bak_partition].img_size), ntohl(pkg_hdr->subimg_hdr[bak_partition].img_crc32));
	if (result) {
		return E_FAILURE;
	}

#ifndef SPIFLASH_XIP_EN
	sfud_xip_dis();
#endif

	return E_OK;
}

static void isp_cmd_get_info_handler(void)
{
	uint8_t i;
	uint8_t size = ARRAY_SIZE(glb_isp_func_set);
	uint8_t len = size + 2;
	uint8_t send_buf[len];

	// bytes to send
	send_buf[0] = len - 1;

	// bootloader version
	send_buf[1] = 0x01;

	//support command list
	for (i = 0; i < size; i++) {
		send_buf[i + 2] = ISP_CMD_FUNC_PTR(i)->cmd_id;
	}

	isp_uart_puts(send_buf, len);

	isp_send_ack();
}

static void isp_cmd_set_baud_handler(void)
{
	uint8_t recv_buf[5];
	uint32_t baud = 0;
	uint8_t cs;

	isp_uart_gets(recv_buf, 5);
	baud = recv_buf[0] << 24;
	baud |= recv_buf[1] << 16;
	baud |= recv_buf[2] << 8;
	baud |= recv_buf[3];
	cs = recv_buf[4];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 4)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	isp_send_ack();

	syscnt_mdelay(1);

	isp_set_uart_baud(baud);
}

static void isp_cmd_read_mem_handler(void)
{
	uint8_t recv_buf[7];
	uint8_t cs;
	uint16_t len = 0;
	uint32_t mem_addr = 0;

	// 4 bytes address + 2 byte len + 1 byte checksum
	isp_uart_gets(recv_buf, 7);
	mem_addr = recv_buf[0] << 24;
	mem_addr |= recv_buf[1] << 16;
	mem_addr |= recv_buf[2] << 8;
	mem_addr |= recv_buf[3];
	len = recv_buf[4] << 8;
	len |= recv_buf[5];
	cs = recv_buf[6];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 6)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	isp_send_ack();

	isp_uart_puts((uint8_t *)mem_addr, len);
}

static void isp_cmd_rst_handler(void)
{
	//TODO
}

static void isp_cmd_write_mem_handler(void)
{
	uint8_t recv_buf[7];
	uint32_t mem_addr = 0;
	uint8_t *addr_ptr = NULL;
	uint8_t cs;
	uint16_t len = 0;
	uint8_t data[ISP_MAX_DATA_LEN];
	uint8_t i;

	// 4 bytes address + 2 byte len + 1 byte checksum
	isp_uart_gets(recv_buf, 7);
	mem_addr = recv_buf[0] << 24;
	mem_addr |= recv_buf[1] << 16;
	mem_addr |= recv_buf[2] << 8;
	mem_addr |= recv_buf[3];
	addr_ptr = (uint8_t *)mem_addr;
	len = recv_buf[4] << 8;
	len |= recv_buf[5];
	cs = recv_buf[6];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 6)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	isp_send_ack();

	// n bytes data + 1 byte checksum
	isp_uart_gets(data, len);
	cs = isp_uart_getc();

	if (E_FAILURE == isp_cmp_checksum(cs, data, len)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	for (i = 0; i < len; i++) {
		addr_ptr[i] = data[i];
	}

	isp_send_ack();
}

#ifdef CONFIG_SDK_QSPIFLASH
static void isp_cmd_read_flash_handler(void)
{
	uint8_t recv_buf[7];
	uint8_t cs;
	uint16_t len = 0;
	uint32_t flash_addr = 0;
	uint8_t data[ISP_MAX_DATA_LEN] = {0};
	sfud_err result = SFUD_SUCCESS;

	// 4 bytes address + 2 byte len + 1 byte checksum
	isp_uart_gets(recv_buf, 7);
	flash_addr = recv_buf[0] << 24;
	flash_addr |= recv_buf[1] << 16;
	flash_addr |= recv_buf[2] << 8;
	flash_addr |= recv_buf[3];
	len = recv_buf[4] << 8;
	len |= recv_buf[5];
	cs = recv_buf[6];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 6)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	result = sfud_read(flash_addr, len, data);
	if (result != SFUD_SUCCESS) {
		isp_send_nack(ISP_READ_FLASH_FAILED);
		return;
	}

	isp_send_ack();

	isp_uart_puts(data, len);
}

static void isp_cmd_write_flash_handler(void)
{
	uint8_t recv_buf[7];
	uint32_t flash_addr = 0;
	uint8_t cs;
	uint16_t len = 0;
	uint8_t data[ISP_MAX_DATA_LEN] = {0};
	sfud_err result = SFUD_SUCCESS;

	// 4 bytes address + 2 byte len + 1 byte checksum
	isp_uart_gets(recv_buf, 7);
	flash_addr = recv_buf[0] << 24;
	flash_addr |= recv_buf[1] << 16;
	flash_addr |= recv_buf[2] << 8;
	flash_addr |= recv_buf[3];
	len = recv_buf[4] << 8;
	len |= recv_buf[5];
	cs = recv_buf[6];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 6)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	result = sfud_erase(flash_addr, len);
	if (result != SFUD_SUCCESS) {
		isp_send_nack(ISP_ERASE_FLASH_FAILED);
		return;
	}

	isp_send_ack();

	// n bytes data + 1 byte checksum
	isp_uart_gets(data, len);
	cs = isp_uart_getc();

	if (E_FAILURE == isp_cmp_checksum(cs, data, len)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	result = sfud_write(flash_addr, len, data);
	if (result != SFUD_SUCCESS) {
		isp_send_nack(ISP_WRITE_FLASH_FAILED);
		return;
	}

	isp_send_ack();
}

static void isp_cmd_erase_flash_handler(void)
{
	uint8_t recv_buf[9];
	uint8_t cs;
	uint32_t len = 0;
	uint32_t flash_addr = 0;
	sfud_err result = SFUD_SUCCESS;

	// 4 bytes address + 2 byte len + 1 byte checksum
	isp_uart_gets(recv_buf, 9);
	flash_addr = recv_buf[0] << 24;
	flash_addr |= recv_buf[1] << 16;
	flash_addr |= recv_buf[2] << 8;
	flash_addr |= recv_buf[3];
	len = recv_buf[4] << 24;
	len |= recv_buf[5] << 16;
	len |= recv_buf[6] << 8;
	len |= recv_buf[7];
	cs = recv_buf[8];

	if (E_FAILURE == isp_cmp_checksum(cs, recv_buf, 8)) {
		isp_send_nack(ISP_CHECKSUM_ERROR);
		return;
	}

	result = sfud_erase(flash_addr, len);

	if (result != SFUD_SUCCESS) {
		isp_send_nack(ISP_ERASE_FLASH_FAILED);
		return;
	}

	isp_send_ack();
}

static int isp_pkg_header_recv(uint32_t boot_partition, uint32_t *part_base, uint32_t *img_size)
{
	uint8_t recv_data[ISP_MAX_DATA_LEN + 2];
	uint16_t data_len = 0;
	uint8_t retrans_num = 0;
	uint32_t flash_addr = ISP_FLASH_PKG_HDR_ADDR;
	uint32_t recv_len = 0;
	uint8_t cs = 0;
	uint32_t i = 0;
	struct pkg_header *pkg_hdr = NULL;
	uint32_t hdr_size = 0;
	sfud_err result = SFUD_SUCCESS;

	do {
		// 2 byts data_len + n bytes data + 1 byte checksum
		isp_uart_gets(recv_data, 2);
		data_len = recv_data[0] << 8;
		data_len |= recv_data[1];

		isp_uart_gets(&recv_data[2], data_len);
		cs = isp_uart_getc();
		if (E_FAILURE == isp_cmp_checksum(cs, recv_data, data_len + 2)) {
			isp_send_nack(ISP_CHECKSUM_ERROR);
			retrans_num++;

			if (retrans_num > ISP_MAX_RETRANS_NUM) {
				return E_FAILURE;
			}

			continue;
		}

		retrans_num = 0;

		if (recv_len == 0) {
			pkg_hdr = (struct pkg_header *)(&recv_data[2]);
			hdr_size = ntohl(pkg_hdr->hdr_size);
			for (i = 0; i < eIMG_MAX_NUM; i++) {
				part_base[i] = ntohl(pkg_hdr->subimg_hdr[i].part_base);
				img_size[i] = ntohl(pkg_hdr->subimg_hdr[i].img_size);
			}

			pkg_hdr->boot_partition = htonl(boot_partition);

			result = sfud_erase(ISP_FLASH_PKG_HDR_ADDR, ISP_FLASH_PKG_HDR_SIZE);
			if (result != SFUD_SUCCESS) {
				isp_send_nack(ISP_ERASE_FLASH_FAILED);
				return E_FAILURE;
			}
		}

		result = sfud_write(flash_addr, data_len, &recv_data[2]);
		if (result != SFUD_SUCCESS) {
			isp_send_nack(ISP_WRITE_FLASH_FAILED);
			return E_FAILURE;
		}

		isp_send_ack();

		flash_addr += data_len;
		recv_len += data_len;
	} while (recv_len < hdr_size);

	return E_OK;
}

static int isp_subimage_recv(uint32_t part_base, uint32_t img_size)
{
	uint8_t recv_data[ISP_MAX_DATA_LEN + 2];
	uint16_t data_len = 0;
	uint8_t retrans_num = 0;
	uint32_t flash_addr = part_base;
	uint32_t recv_len = 0;
	uint8_t cs = 0;
	sfud_err result = SFUD_SUCCESS;

	do {
		// 2 byts data_len + n bytes data + 1 byte checksum
		isp_uart_gets(recv_data, 2);
		data_len = recv_data[0] << 8;
		data_len |= recv_data[1];

		isp_uart_gets(&recv_data[2], data_len);
		cs = isp_uart_getc();
		if (E_FAILURE == isp_cmp_checksum(cs, recv_data, data_len + 2)) {
			isp_send_nack(ISP_CHECKSUM_ERROR);
			retrans_num++;

			if (retrans_num > ISP_MAX_RETRANS_NUM) {
				return E_FAILURE;
			}

			continue;
		}

		retrans_num = 0;

		if (recv_len == 0) {
			result = sfud_erase(part_base, img_size);
			if (result != SFUD_SUCCESS) {
				isp_send_nack(ISP_ERASE_FLASH_FAILED);
				return E_FAILURE;
			}
		}

		result = sfud_write(flash_addr, data_len, &recv_data[2]);
		if (result != SFUD_SUCCESS) {
			isp_send_nack(ISP_WRITE_FLASH_FAILED);
			return E_FAILURE;
		}

		isp_send_ack();

		flash_addr += data_len;
		recv_len += data_len;
	} while (recv_len < img_size);

	return E_OK;
}

static uint32_t isp_get_boot_partition(void)
{
	uint8_t recv_buf[4];
	uint32_t boot_partition;
	sfud_err result = SFUD_SUCCESS;
	uint32_t offset;

	offset = offsetof(struct pkg_header, boot_partition);
	result = sfud_read(ISP_FLASH_PKG_HDR_ADDR + offset, 4, recv_buf);
	if (result != SFUD_SUCCESS) {
		return eBOOT_PARTITION_UNKNOWN;
	}

	boot_partition = ntohl(*(uint32_t *)recv_buf);

	return boot_partition;
}

static uint32_t isp_switch_boot_partition(void)
{
	sfud_err result = SFUD_SUCCESS;
	struct pkg_header *pkg_hdr = NULL;
	uint32_t boot_partition;

	result = sfud_read(ISP_FLASH_PKG_HDR_ADDR, ISP_FLASH_PKG_HDR_SIZE, g_flash_tmp_buf);
	if (result != SFUD_SUCCESS) {
		return ISP_READ_FLASH_FAILED;
	}

	pkg_hdr = (struct pkg_header *)g_flash_tmp_buf;

	boot_partition = ntohl(pkg_hdr->boot_partition);
	if (boot_partition == eBOOT_PARTITION1)
		pkg_hdr->boot_partition = htonl(eBOOT_PARTITION2);
	else
		pkg_hdr->boot_partition = htonl(eBOOT_PARTITION1);

	result = sfud_erase(ISP_FLASH_PKG_HDR_ADDR, ISP_FLASH_PKG_HDR_SIZE);
	if (result != SFUD_SUCCESS) {
		return ISP_ERASE_FLASH_FAILED;
	}

	result = sfud_write(ISP_FLASH_PKG_HDR_ADDR, ISP_FLASH_PKG_HDR_SIZE, g_flash_tmp_buf);
	if (result != SFUD_SUCCESS) {
		return ISP_WRITE_FLASH_FAILED;
	}

	return 0;
}

static void isp_cmd_upgrade_handler(void)
{
	uint8_t i = 0;
	uint32_t part_base[eIMG_MAX_NUM] = {0};
	uint32_t img_size[eIMG_MAX_NUM] = {0};
	uint32_t boot_partition;
	uint32_t bak_partition;

	boot_partition = isp_get_boot_partition();
	if (boot_partition == eBOOT_PARTITION_UNKNOWN) {
		isp_send_nack(ISP_READ_FLASH_FAILED);
		return ;
	} else if (boot_partition == eBOOT_PARTITION2) {
		bak_partition = eBOOT_PARTITION1;
	}
	else {
		bak_partition = eBOOT_PARTITION2;
	}

	if (E_OK != isp_pkg_header_recv(boot_partition, part_base, img_size))
		return ;

	for (i = 0; i < eIMG_MAX_NUM; i++) {
		if (E_OK !=isp_subimage_recv(part_base[i], img_size[i])) {
			return ;
		}
	}

	if (E_OK != isp_upgrade_pkg_check(bak_partition)) {
		isp_send_nack(ISP_IMAGE_CHECK_FAILED);
		return ;
	}

	isp_send_ack();
}

static void isp_cmd_get_boot_part_handler(void)
{
	uint32_t boot_partition = isp_get_boot_partition();

	if (boot_partition == eBOOT_PARTITION_UNKNOWN) {
		isp_send_nack(ISP_READ_FLASH_FAILED);
		return ;
	}

	isp_uart_putc((uint8_t)boot_partition);

	isp_send_ack();
}

static void isp_cmd_switch_boot_part_handler(void)
{
	uint32_t result = isp_switch_boot_partition();
	if (result) {
		isp_send_nack(result);
		return ;
	}

	isp_send_ack();
}

#endif

static int32_t isp_cmd_check(uint8_t cmd, uint8_t *idx)
{
	uint8_t i;
	uint8_t size = ARRAY_SIZE(glb_isp_func_set);

	for (i = 0; i < size; i++) {
		if (ISP_CMD_FUNC_PTR(i)->cmd_id == cmd) {
			*idx = i;
			return E_OK;
		}
	}

	return E_FAILURE;
}

static void isp_sync(void)
{
	uint8_t sync;

	while (1) {
		sync = isp_uart_getc();

		if (sync == ISP_SYNC_DATA)
			break;
	}

	isp_send_ack();
}

void isp_start(void)
{
	uint8_t cmd;
	uint8_t chksum;
	uint8_t idx;

	printf_pause();

	isp_sync();

	while (1) {
		cmd = isp_uart_getc();
		if (E_OK != isp_cmd_check(cmd, &idx)) {
			isp_send_nack(ISP_UNSUPPORTED_CMD);
			continue;
		}

		chksum = isp_uart_getc();
		if (0xFF != cmd + chksum) {
			isp_send_nack(ISP_CHECKSUM_ERROR);
			continue;
		}

		isp_send_ack();

		if (cmd == ISP_CMD_STOP)
			break;

		ISP_CMD_FUNC_PTR(idx)->func();
	}

	printf_restart();
}

